# -*- coding:utf-8 -*-
# @FileName : app\services\models\yuanjing_model.py
# @Time     : 2025/11/28
# @Author   : 天空之城
"""模拟OpenAI的Function对象"""
# yuanjing_model.py
import re
import requests
import datetime
import hashlib
from typing import List, Optional, Dict, Any
from app.config import (
    logger,
    YUANJING_APP_ID,
    YUANJING_NLPT_AUTHORIZATION,
    YUANJING_APP_SECRET,
    YUANJING_MODEL,
    YUANJING_BASE_URL,
    YUANJING_MAX_RETRIES
)


class FunctionWrapper:
    
    
    def __init__(self, function_dict: Dict[str, Any]):
        """
        初始化Function对象
        
        Args:
            function_dict: 包含function信息的字典
        """
        self.name = function_dict.get("name", "")
        self.arguments = function_dict.get("arguments", "{}")
    
    def __repr__(self) -> str:
        """返回字典格式的字符串表示"""
        return f"{{'name': {repr(self.name)}, 'arguments': {repr(self.arguments)}}}"
    
    def __str__(self) -> str:
        """返回字典格式的字符串表示"""
        return self.__repr__()


class ToolCallWrapper:
    """模拟OpenAI的ToolCall对象"""
    
    def __init__(self, tool_call_dict: Dict[str, Any]):
        """
        初始化ToolCall对象
        
        Args:
            tool_call_dict: 包含tool_call信息的字典
        """
        self.id = tool_call_dict.get("id", "")
        self.type = tool_call_dict.get("type", "function")
        
        # 处理function字段
        function_data = tool_call_dict.get("function", {})
        if isinstance(function_data, dict):
            self.function = FunctionWrapper(function_data)
        else:
            self.function = function_data
    
    def __repr__(self) -> str:
        """返回字典格式的字符串表示"""
        func_repr = repr(self.function) if hasattr(self.function, '__repr__') else str(self.function)
        return f"{{'id': {repr(self.id)}, 'type': {repr(self.type)}, 'function': {func_repr}}}"
    
    def __str__(self) -> str:
        """返回字典格式的字符串表示"""
        return self.__repr__()


class ChoiceMessage:
    """模拟OpenAI的ChatCompletionMessage对象"""
    
    def __init__(self, message_dict: Dict[str, Any]):
        """
        初始化消息对象
        
        Args:
            message_dict: 包含message信息的字典
        """
        self.role = message_dict.get("role", "assistant")
        self.content = message_dict.get("content", "")
        self.tool_call_id = message_dict.get("tool_call_id")
        
        # 处理tool_calls字段，转换为对象列表
        tool_calls_data = message_dict.get("tool_calls")
        if tool_calls_data:
            if isinstance(tool_calls_data, list):
                self.tool_calls = [
                    ToolCallWrapper(tc) if isinstance(tc, dict) else tc 
                    for tc in tool_calls_data
                ]
            else:
                self.tool_calls = tool_calls_data
        else:
            self.tool_calls = None
    
    def model_dump(self) -> Dict[str, Any]:
        """
        返回字典格式，兼容Pydantic的model_dump方法
        
        Returns:
            消息的字典表示
        """
        result = {
            "role": self.role,
            "content": self.content
        }
        if self.tool_calls is not None:
            # 将tool_calls对象列表转换为字典列表
            if isinstance(self.tool_calls, list):
                result["tool_calls"] = [
                    {
                        "id": tc.id if hasattr(tc, 'id') else "",
                        "type": tc.type if hasattr(tc, 'type') else "function",
                        "function": {
                            "name": tc.function.name if hasattr(tc, 'function') else "",
                            "arguments": tc.function.arguments if hasattr(tc, 'function') else "{}"
                        }
                    } if hasattr(tc, 'function') else tc
                    for tc in self.tool_calls
                ]
            else:
                result["tool_calls"] = self.tool_calls
        if self.tool_call_id is not None:
            result["tool_call_id"] = self.tool_call_id
        return result
    
    def __repr__(self) -> str:
        """返回字典格式的字符串表示，与之前打印字典格式一致"""
        result = {
            "role": self.role,
            "content": self.content
        }
        if self.tool_calls is not None:
            if isinstance(self.tool_calls, list):
                result["tool_calls"] = [
                    {
                        "id": tc.id if hasattr(tc, 'id') else "",
                        "type": tc.type if hasattr(tc, 'type') else "function",
                        "function": {
                            "name": tc.function.name if hasattr(tc, 'function') else "",
                            "arguments": tc.function.arguments if hasattr(tc, 'function') else "{}"
                        }
                    } if hasattr(tc, 'function') else tc
                    for tc in self.tool_calls
                ]
            else:
                result["tool_calls"] = self.tool_calls
        else:
            result["tool_calls"] = None
        if self.tool_call_id is not None:
            result["tool_call_id"] = self.tool_call_id
        return repr(result)
    
    def __str__(self) -> str:
        """返回字典格式的字符串表示"""
        return self.__repr__()


class ChoiceWrapper:
    """模拟OpenAI的Choice对象，将字典转换为对象"""
    
    def __init__(self, choice_dict: Dict[str, Any]):
        """
        初始化Choice对象
        
        Args:
            choice_dict: 包含choice信息的字典
        """
        self.index = choice_dict.get("index", 0)
        self.finish_reason = choice_dict.get("finish_reason", "stop")
        
        # 处理message字段
        message_data = choice_dict.get("message", {})
        if isinstance(message_data, dict):
            self.message = ChoiceMessage(message_data)
        else:
            # 如果已经是对象，直接使用
            self.message = message_data
    
    def __repr__(self) -> str:
        """返回字典格式的字符串表示"""
        message_repr = repr(self.message) if hasattr(self.message, '__repr__') else str(self.message)
        return f"{{'index': {self.index}, 'finish_reason': {repr(self.finish_reason)}, 'message': {message_repr}}}"
    
    def __str__(self) -> str:
        """返回字典格式的字符串表示"""
        return self.__repr__()


from app.services.models.base import BaseModel

class YuanjingModel(BaseModel):
    """元景大模型调用封装"""
    
    def __init__(
        self, 
        app_id: str = YUANJING_APP_ID,
        nlpt_authorization: str = YUANJING_NLPT_AUTHORIZATION,
        app_secret: str = YUANJING_APP_SECRET,
        model: str = YUANJING_MODEL,
        base_url: str = YUANJING_BASE_URL,
        max_retries: int = None
    ):
        """
        初始化元景大模型
        
        Args:
            app_id: 接入标识码
            nlpt_authorization: nlpt-Authorization的Bearer token
            app_secret: 用于计算TOKEN的密钥
            model: 使用的模型名称，默认为 Qwen2.5-32B
            base_url: API基础URL
            max_retries: 最大重试次数
        """
        if max_retries is None:
            max_retries = YUANJING_MAX_RETRIES
        self.app_id = app_id
        self.nlpt_authorization = nlpt_authorization
        self.app_secret = app_secret
        self.model = model
        self.base_url = base_url
        self.max_retries = max_retries
        self.logger = logger
        
        # 默认参数配置
        self.default_params = {
            "temperature": 0.1,
            "max_tokens": 2000,
            "stream": False,
            "presence_penalty": 1.03,
            "frequency_penalty": 0.0,
            "seed": None,
            "top_p": 0.95,
        }
    
    def _generate_timestamp(self) -> str:
        """生成时间戳，格式：yyyy-MM-dd hh:mm:ss SSS"""
        return str(datetime.datetime.now())[:-3].replace('.', ' ')
    
    def _generate_trans_id(self, timestamp: str) -> str:
        """生成序列号，格式：yyyyMMddhhmmssSSS+6位随机数"""
        trans_id = timestamp.replace('-', '').replace(':', '').replace(' ', '').strip() + '123456'
        return trans_id
    
    def _calculate_token(self, timestamp: str, trans_id: str) -> str:
        """
        计算TOKEN值
        根据系统参数名称（除token）将所有请求系统参数按照字母先后顺序排序
        然后拼接app_secret，最后进行MD5(32位小写)加密
        """
        # 按字母顺序排序参数：APP_ID, TIMESTAMP, TRANS_ID
        sb = ["APP_ID", self.app_id, "TIMESTAMP", timestamp, "TRANS_ID", trans_id, self.app_secret]
        codestring = ''.join(sb)
        
        # MD5的32位小写
        tokener = hashlib.md5()
        tokener.update(codestring.encode('utf-8'))
        token = tokener.hexdigest().lower()
        
        return token
    
    def _build_request_body(
        self, 
        messages: List[Dict[str, str]], 
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: str = "auto",
        **kwargs
    ) -> Dict[str, Any]:
        """
        构建请求体
        
        Args:
            messages: 对话消息列表
            tools: 工具定义列表
            tool_choice: 工具选择策略
            **kwargs: 其他参数（temperature, max_tokens等）
        
        Returns:
            完整的请求体
        """
        # 生成时间戳和序列号
        timestamp = self._generate_timestamp()
        trans_id = self._generate_trans_id(timestamp)
        
        # 计算TOKEN
        token = self._calculate_token(timestamp, trans_id)
        
        # 合并参数
        params = {**self.default_params, **kwargs}
        
        # 构建请求体
        req_body_params = {
            "model": self.model,
            "messages": messages,
            "temperature": params["temperature"],
            "max_tokens": params["max_tokens"],
            "stream": params["stream"],
            "presence_penalty": params["presence_penalty"],
            "frequency_penalty": params["frequency_penalty"],
            "seed": params["seed"],
            "top_p": params["top_p"],
        }
        
        # 如果params中有top_k，则添加
        if "top_k" in params:
            req_body_params["top_k"] = params["top_k"]
        
        req_data = {
            "UNI_BSS_HEAD": {
                "APP_ID": self.app_id,
                "TIMESTAMP": timestamp,
                "TRANS_ID": trans_id,
                "TOKEN": token
            },
            "UNI_BSS_BODY": {
                "YUANJING_MODEL_REQ": req_body_params
            },
            "UNI_BSS_ATTACHED": {
                "MEDIA_INFO": ""
            }
        }
        
        # 如果提供了tools，添加到请求中
        if tools:
            req_data["UNI_BSS_BODY"]["YUANJING_MODEL_REQ"]["tools"] = tools
            req_data["UNI_BSS_BODY"]["YUANJING_MODEL_REQ"]["tool_choice"] = tool_choice
        
        return req_data
    
    def _build_headers(self) -> Dict[str, str]:
        """构建请求头"""
        return {
            "Accept-Encoding": "",
            "content-type": "application/json; charset=UTF-8",
            "Accept": "application/json",
            "nlpt-Authorization": f"Bearer {self.nlpt_authorization}",
            "scene_code": "test"
        }
    
    # def _extract_choice(self, response_data: Dict[str, Any]) -> Dict[str, Any]:
    #     """
    #     从响应中提取choices[0]，与moda保持一致
        
    #     Args:
    #         response_data: API响应数据
        
    #     Returns:
    #         choices[0]对象
    #     """
    #     try:
    #         # 实际的API返回格式：直接在顶层有choices字段（OpenAI格式）
    #         if "choices" in response_data and len(response_data["choices"]) > 0:
    #             return response_data["choices"][0]
            
    #         # 兼容UNI_BSS_BODY格式（如果存在）
    #         if "UNI_BSS_BODY" in response_data:
    #             body = response_data["UNI_BSS_BODY"]
    #             if "YUANJING_MODEL_RSP" in body:
    #                 rsp = body["YUANJING_MODEL_RSP"]
    #                 if "choices" in rsp and len(rsp["choices"]) > 0:
    #                     return rsp["choices"][0]
            
    #         # 如果结构不同，抛出异常
    #         raise ValueError("响应结构不符合预期，未找到choices字段")
            
    #     except Exception as e:
    #         self.logger.error(f"提取响应内容失败: {str(e)}")
    #         raise Exception(f"提取响应内容失败: {str(e)}") from e
    
    def generate_response(
        self, 
        messages: List[Dict[str, str]], 
        tools: Optional[List[Dict[str, Any]]] = None,
        tool_choice: str = "auto",
        temperature: float = 0.1,
        max_tokens: int = 2000,
        **kwargs
    ) -> Dict[str, Any]:
        """
        调用元景大模型API生成响应
        
        Args:
            messages: 对话消息列表
            tools: 工具定义列表（可选）
            tool_choice: 工具选择策略，默认为auto
            temperature: 生成温度，控制随机性 (0-1)，默认0.1
            max_tokens: 最大生成token数，默认2000
            **kwargs: 其他参数（stream, presence_penalty等）
        
        Returns:
            choices[0]对象，与moda保持一致
        
        Raises:
            Exception: API调用失败时抛出
        """
        # 合并参数
        params = {
            "temperature": temperature,
            "max_tokens": max_tokens,
            **kwargs
        }
        
        # 构建请求
        req_data = self._build_request_body(messages, tools, tool_choice, **params)
        headers = self._build_headers()
        
        # 重试机制
        last_error = None
        for retry in range(self.max_retries + 1):
            try:
                self.logger.info(f"[元景模型] 调用API，模型: {self.model}，第 {retry + 1}/{self.max_retries + 1} 次尝试")
                
                response = requests.post(
                    self.base_url,
                    json=req_data,
                    headers=headers,
                    stream=False,
                    timeout=120
                )
                
                # 检查HTTP状态码
                response.raise_for_status()
                
                # 解析响应
                response_data = response.json()
                
                # 检查响应码
                if "UNI_BSS_HEAD" in response_data:
                    resp_code = response_data["UNI_BSS_HEAD"].get("RESP_CODE", "")
                    resp_desc = response_data["UNI_BSS_HEAD"].get("RESP_DESC", "")
                    
                    if resp_code != "00000":
                        error_msg = f"API返回错误: RESP_CODE={resp_code}, RESP_DESC={resp_desc}"
                        self.logger.error(f"[元景模型] {error_msg}")
                        raise Exception(error_msg)
                
                # 提取choices[0]，转换为对象格式以兼容OpenAI的Choice接口
                choice_dict = response_data.get("choices", [{}])[0]
                
                # 将字典转换为ChoiceWrapper对象，使其与MoDaModel返回的格式一致
                choice = ChoiceWrapper(choice_dict)
                self.logger.info(f"[元景模型] API调用成功，模型: {self.model}")
                return choice
                
            except requests.exceptions.RequestException as e:
                last_error = f"网络请求失败: {str(e)}"
                self.logger.warning(f"[元景模型] {last_error}，重试中 ({retry+1}/{self.max_retries+1})")
                
            except Exception as e:
                last_error = f"API调用失败: {str(e)}"
                self.logger.error(f"[元景模型] {last_error}，重试中 ({retry+1}/{self.max_retries+1})")
            
            if retry < self.max_retries:
                import time
                time.sleep(1)  # 等待1秒后重试
        
        # 所有重试都失败
        error_msg = f"所有重试均失败，最后一次错误: {last_error}"
        self.logger.error(f"[元景模型] {error_msg}")
        raise Exception(error_msg)

